#include "PLLinearEquationSystem.h"

#if PL_LINEAR_EQUATION_SYSTEM_SOLVER == PL_SOLVER_IML
	#include "diagpre_double.h"              // Diagonal preconditioner
	#include "icpre_double.h"
	#include "cg.h"                          // IML++ CG template
	#include "bicgstab.h"                    // IML++ BiCGSTAB template
#endif

#include <iostream>

using std::cout;
using std::endl;
using std::map;
using std::vector;

/**
  *	Constructs a new linear equation system of the form A*x=b.
  * The matrix A gets the dimensions numRows x numCols
  *
  * @param numRows number of rows of the matrix A in the equation system A*x=b
  * @param numCols number of columns of the matrix A in the equation system A*x=b
  */
PLLinearEquationSystem::PLLinearEquationSystem( const PLuint &numRows, const PLuint &numCols )
{
	mNumRows = numRows;
	mNumCols = numCols;

#if PL_LINEAR_EQUATION_SYSTEM_SOLVER == PL_SOLVER_IML
	mMatrixA = 0;
	mVectorB = 0;
	mVectorX = 0;
#endif

#if PL_LINEAR_EQUATION_SYSTEM_SOLVER == PL_SOLVER_OPENCV
	mMatrixA = 0;
	mVectorB = 0;
	mVectorX = 0;
#endif

	Initialize();
}

/**
  *	Destroys the PLLinearEquationSystem object. All data is deleted.
  */
PLLinearEquationSystem::~PLLinearEquationSystem( void )
{
#if PL_LINEAR_EQUATION_SYSTEM_SOLVER == PL_SOLVER_IML
	if (mMatrixA) delete mMatrixA;
	if (mVectorB) delete mVectorB;
	if (mVectorX) delete mVectorX;
#endif

#if PL_LINEAR_EQUATION_SYSTEM_SOLVER == PL_SOLVER_OPENCV
	if (mMatrixA) cvReleaseMat( &mMatrixA );
	if (mVectorX) cvReleaseMat( &mVectorX );
	if (mVectorB) cvReleaseMat( &mVectorB );
#endif
}

/**
  *	Fills the matrix A. Matrix A must be of the form JT*J, where JT is the transpose
  * of J. This method must be called line per line for the matrix J, because
  *	it fills the matrix JTJ (=A) for a single row of J.
  *
  * @param variables indexes of the variables in the current row of the matrix J
  * @param coefficients values of the variables in the row of J, which means that in the current row of J the variable at index variables[i] has the value coefficients[i]
  * @param weight weight of the variables
  * @param f function value corresponding to the current row of J
  */
void PLLinearEquationSystem::FillA( const vector<PLuint> &variables,
	const vector<double> &coefficients, const double &weight, const double &f )
{
	for (PLuint i = 0; i < variables.size(); i++)
	{
		for (PLuint j = 0; j < variables.size(); j++)
		{

#if PL_LINEAR_EQUATION_SYSTEM_SOLVER == PL_SOLVER_IML
			mMatrixCoordinateToValueMap[PLMatrixCoordinate( variables[i], variables[j] )] += coefficients[i] * coefficients[j] * weight * weight;
#endif

#if PL_LINEAR_EQUATION_SYSTEM_SOLVER == PL_SOLVER_OPENCV
			CV_MAT_ELEM( *mMatrixA, double, variables[i], variables[j] ) += coefficients[i] * coefficients[j] * weight * weight;
#endif
		}

#if PL_LINEAR_EQUATION_SYSTEM_SOLVER == PL_SOLVER_IML
		mVectorIndexToValueMap[ variables[i] ] -= weight * weight * f * coefficients[i];
#endif

#if PL_LINEAR_EQUATION_SYSTEM_SOLVER == PL_SOLVER_OPENCV
		CV_MAT_ELEM( *mVectorB, double, variables[i], 0 ) -= weight * weight * f * coefficients[i];
#endif
		
	}
}

/**
  *	Must be called after the linear equation system is built and is ready for solving.
  */
void PLLinearEquationSystem::Finalize( void )
{
#if PL_LINEAR_EQUATION_SYSTEM_SOLVER == PL_SOLVER_IML
	if (mMatrixA) delete mMatrixA;
	if (mVectorB) delete mVectorB;
	if (mVectorX) delete mVectorX;

	int M = static_cast<int> ( mNumRows );
	int N = static_cast<int> ( mNumCols );
	int nz = static_cast<int>( mMatrixCoordinateToValueMap.size() );
	
	vector<double> val;
	val.reserve( nz );

	vector<int> r;
	
	vector<int> c;
	c.reserve( nz );

	vector<double> b( M, 0.0 );

	int current_val_index = 0;
	int current_row_index = -1;
	
	for (map<PLMatrixCoordinate, double, PLMatrixCoordinate::PLMatrixCoordinateCompare>::iterator it = mMatrixCoordinateToValueMap.begin();
		it != mMatrixCoordinateToValueMap.end(); it++)
	{
		val.push_back( it->second );

		c.push_back( static_cast<int>( it->first.Col() ) );

		if (current_row_index != it->first.Row())
		{
			// New row
			r.push_back( current_val_index );

			current_row_index = static_cast<int>( it->first.Row() );
		}

		current_val_index++;
	}

	// Number of non-zeros at the end of the r-array
	r.push_back( nz );

	for (map<PLuint, double>::iterator it = mVectorIndexToValueMap.begin();
		it != mVectorIndexToValueMap.end(); it++)
	{
		b[it->first] = it->second;
	}

	mMatrixA = new CompRow_Mat_double( M, N, nz, &val[0], &r[0], &c[0] );
	mVectorB = new VECTOR_double( &b[0], M );
	mVectorX = new VECTOR_double( N, 0.0 );

#endif
}

/**
  *	Returns the number of rows of the matrix A.
  *
  * @return number of rows of the matrix A
  */
const PLuint &PLLinearEquationSystem::GetNumRows( void ) const
{
	return mNumRows;
}

/**
  *	Returns the number of columns of the matrix A.
  *
  * @return number of columns of the matrix A
  */
const PLuint &PLLinearEquationSystem::GetNumCols( void ) const
{
	return mNumCols;
}

/**
  *	Initializes the linear equation system
  */
void PLLinearEquationSystem::Initialize( void )
{
#if PL_LINEAR_EQUATION_SYSTEM_SOLVER == PL_SOLVER_OPENCV
	if (mMatrixA) cvReleaseMat( &mMatrixA );
	if (mVectorX) cvReleaseMat( &mVectorX );
	if (mVectorB) cvReleaseMat( &mVectorB );

	mMatrixA = cvCreateMat( static_cast<int>(mNumRows),  static_cast<int>(mNumCols), CV_64FC1 );
	mVectorB = cvCreateMat( static_cast<int>(mNumRows), 1, CV_64FC1 );
	mVectorX = cvCreateMat( static_cast<int>(mNumCols), 1, CV_64FC1 );

	Reset();
#endif
}

/**
  *	Resets the matrix A and vectors x and b to contain just 0-values
  */
void PLLinearEquationSystem::Reset( void )
{
#if PL_LINEAR_EQUATION_SYSTEM_SOLVER == PL_SOLVER_IML
	for (map<PLMatrixCoordinate, double, PLMatrixCoordinate::PLMatrixCoordinateCompare>::iterator it = mMatrixCoordinateToValueMap.begin();
		it != mMatrixCoordinateToValueMap.end(); it++)
	{
		it->second = 0.0;
	}

	for (map<PLuint, double>::iterator it = mVectorIndexToValueMap.begin();
		it != mVectorIndexToValueMap.end(); it++)
	{
		it->second = 0.0;
	}
#endif

#if PL_LINEAR_EQUATION_SYSTEM_SOLVER == PL_SOLVER_OPENCV
	for (PLuint i = 0; i < mNumRows; i++)
	{
		for (PLuint j = 0; j < mNumCols; j++)
		{
			CV_MAT_ELEM( *mMatrixA, double, i, j ) = 0.0;
		}

		CV_MAT_ELEM( *mVectorB, double, i, 0 ) = 0.0;
		CV_MAT_ELEM( *mVectorX, double, i, 0 ) = 0.0;
	}
#endif
}

/**
  *	Solves the linear equation system. The solution vector x can be fetched
  * via the X() method.
  */
void PLLinearEquationSystem::Solve( void )
{
#if PL_LINEAR_EQUATION_SYSTEM_SOLVER == PL_SOLVER_IML
	double tol = 1.e-6;                    // Convergence tolerance
	int maxit = 1500;               // Maximum iterations

	ICPreconditioner_double D( *mMatrixA );	// Create diagonal preconditioner
	CG( *mMatrixA, *mVectorX, *mVectorB, D, maxit, tol );   // Solve system

	cout << "Iterations used: " << maxit << endl;
	cout << "Residual in linear equation system: " << tol << endl;
#endif
#if PL_LINEAR_EQUATION_SYSTEM_SOLVER == PL_SOLVER_OPENCV
	cvSolve( mMatrixA, mVectorB, mVectorX );
#endif
}

/**
  *	Returns the value of the solution vector, that corresponds to the passed index
  *
  * @param index index of the value in the solution vector
  * @return value of the solution vector, that corresponds to the passed index
  */
const double &PLLinearEquationSystem::X( const PLuint index ) const
{
#if PL_LINEAR_EQUATION_SYSTEM_SOLVER == PL_SOLVER_IML
	return (*mVectorX)( static_cast<int>( index ) );
#endif

#if PL_LINEAR_EQUATION_SYSTEM_SOLVER == PL_SOLVER_OPENCV
	return CV_MAT_ELEM( *mVectorX, double, index, 0 );
#endif
}